// File:          webots_controller.cpp
// Date:
// Description:
// Author:
// Modifications:

#include "webots_controller.h"
#include "common/cppTypes.h"
#include "common/Math/orientation_tools.h"
#include "common/TouchForce.h"

using namespace webots;

std::string motor_name[4][3] = { "RF-joint-1", "RF-joint-2", "RF-joint-3", "LF-joint-1", "LF-joint-2", "LF-joint-3",
                                 "RB-joint-1", "RB-joint-2", "RB-joint-3", "LB-joint-1", "LB-joint-2", "LB-joint-3" };

std::string sensor_name[4][3] = { "RF-joint-1_sensor", "RF-joint-2_sensor", "RF-joint-3_sensor", "LF-joint-1_sensor",
                                  "LF-joint-2_sensor", "LF-joint-3_sensor", "RB-joint-1_sensor", "RB-joint-2_sensor",
                                  "RB-joint-3_sensor", "LB-joint-1_sensor", "LB-joint-2_sensor", "LB-joint-3_sensor" };
                            
std::string tc_sensors[4] = {"RF_touch", "LF_touch", "RB_touch", "LB_touch"};

void webotsController::initNode(void) {
    legdata_puber = _n->advertise<common::LegData>("/quadruped/control/leg_data", 10);
    touchstate_puber = _n->advertise<common::TouchForce>("/quadruped/control/touch_force", 10);
    state_puber = _n->advertise<common::CheaterState>("/quadruped/control/cheater_state", 10);
    imu_puber = _n->advertise<common::ImuData>("/quadruped/control/imu_data", 10);
    cmd_suber = _n->subscribe("/quadruped/control/command", 100, &webotsController::robotCmd_Callback, this);
    initDevice();
}

/**
 * @brief Initiate and enable devices
 *
 */
void webotsController::initDevice(void) {
    imu = robot->getInertialUnit("imu");
    imu->enable(timeStep);

    accel = robot->getAccelerometer("accelerometer");
    accel->enable(timeStep);

    gyro = robot->getGyro("gyro");
    gyro->enable(timeStep);

    for (int i = 0; i < 12; i++) {
        m[i / 3][i % 3] = robot->getMotor(motor_name[i / 3][i % 3]);
        p[i / 3][i % 3] = robot->getPositionSensor(sensor_name[i / 3][i % 3]);
        p[i / 3][i % 3]->enable(timeStep);

        // 使用supervisor来观测电机转速
        motor_v_visor[i / 3][i % 3] = robot->getFromDef(motor_name[i / 3][i % 3]);
    }

    // for(int i = 0; i < 4; i ++){
    //     tc[i] = robot->getTouchSensor(tc_sensors[i]);
    //     tc[i]->enable(timeStep);
    // }
    
    // 由于模型的命名冲突，这两个关节的DEF与关节名称不同
    motor_v_visor[2][1] = robot->getFromDef("RB2");
    motor_v_visor[3][1] = robot->getFromDef("LB2");

    // 初始化世界坐标系下的pv观测器
    world_p_v_visor = robot->getFromDef("Zouwu");
}

void webotsController::updateLegData(void) {
    const double* vel;

    for (int leg = 0; leg < 4; leg++) {
        leg_data.q_abad[leg] = p[leg][0]->getValue();
        leg_data.q_hip[leg] = p[leg][1]->getValue() + ang_bias;
        leg_data.q_knee[leg] = - p[leg][2]->getValue() - ang_bias;     

        vel = motor_v_visor[leg][0]->getVelocity();  leg_data.qd_abad[leg]  = vel[5];
        vel = motor_v_visor[leg][1]->getVelocity();  leg_data.qd_hip[leg]  = -vel[5];
        vel = motor_v_visor[leg][2]->getVelocity();  leg_data.qd_knee[leg]  = -vel[5];
    }
    legdata_puber.publish(leg_data);
}

void webotsController::updateCheaterState(void) {
    const double* data;
    // orientation
    data = imu->getRollPitchYaw();
    Vec3<float> rpy(-data[0], data[1], data[2]);
    Quat<float> q = ori::rpyToQuat(rpy);
    cheater_state.bodyOrientation.w = q.w();
    cheater_state.bodyOrientation.x = q.x();
    cheater_state.bodyOrientation.y = q.y();
    cheater_state.bodyOrientation.z = q.z();
    
    // acceleration
    data = accel->getValues();
    cheater_state.acceleration[0] = -data[0]; 
    cheater_state.acceleration[1] = data[1];
    cheater_state.acceleration[2] = data[2];

    // position
    data = world_p_v_visor->getPosition();
    cheater_state.position[0] = -data[0];
    cheater_state.position[1] = data[2];
    cheater_state.position[2] = data[1] + 0.3633;

    // linear velocity
    data = world_p_v_visor->getVelocity();
    cheater_state.vBody[0] = -data[0];
    cheater_state.vBody[1] = data[2];
    cheater_state.vBody[2] = data[1];

    // angular velocity
    cheater_state.omegaBody[0] = -data[3];
    cheater_state.omegaBody[1] = data[5];
    cheater_state.omegaBody[2] = data[4];

    // publish
    state_puber.publish(cheater_state);
}

void webotsController::updateImuData(void){
    const double* data;

    // quat
    data = imu->getRollPitchYaw();
    Vec3<float> rpy(-data[0], data[1], data[2]);
    Quat<float> q = ori::rpyToQuat(rpy);
    imu_data.quat.w = q.w();
    imu_data.quat.x = q.x();
    imu_data.quat.y = q.y();
    imu_data.quat.z = q.z();

    // accel
    data = accel->getValues();
    imu_data.accel[0] = data[0];
    imu_data.accel[1] = data[1];
    imu_data.accel[2] = data[2];

    // gyro
    data = gyro->getValues();
    imu_data.gyro[0] = data[0];
    imu_data.gyro[1] = data[1];
    imu_data.gyro[2] = data[2];

    // publish
    imu_puber.publish(imu_data);
}

void webotsController::updateTouchForce(void){
    common::TouchForce tc_force;
    const double* data;
    for (int leg = 0; leg < 4; leg++) {
        data = tc[leg]->getValues();
        tc_force.force[leg].x = data[0];
        tc_force.force[leg].y = data[1];
        tc_force.force[leg].z = data[2];
    }
    touchstate_puber.publish(tc_force);
}

void webotsController::robotCmd_Callback(const common::RobotCommandConstPtr& msg) {
    float tau;
    static int64_t iter = 0;
    // location control & force control
    for (int leg = 0; leg < 4; leg++) {
        if(msg->flags[leg]){
            // abad
            if(msg->mode_abad[leg]){
                tau = msg->kp_abad[leg] * (msg->q_des_abad[leg] - leg_data.q_abad[leg]) +
                    msg->kd_abad[leg] * (msg->qd_des_abad[leg] - leg_data.qd_abad[leg]) + msg->tau_abad_ff[leg];
                m[leg][0]->setTorque(tau);
            }
            else m[leg][0]->setPosition(msg->q_des_abad[leg]);

            // hip
            if(msg->mode_hip[leg]){
                tau = msg->kp_hip[leg] * (msg->q_des_hip[leg] - leg_data.q_hip[leg]) +
                    msg->kd_hip[leg] * (msg->qd_des_hip[leg] - leg_data.qd_hip[leg]) + msg->tau_hip_ff[leg];
                m[leg][1]->setTorque(tau);
            }
            else m[leg][1]->setPosition(msg->q_des_hip[leg] - ang_bias);

            // knee
            if(msg->mode_hip[leg]){
                tau = msg->kp_knee[leg] * (msg->q_des_knee[leg] - leg_data.q_knee[leg]) +
                    msg->kd_knee[leg] * (msg->qd_des_knee[leg] - leg_data.qd_knee[leg]) + msg->tau_knee_ff[leg];
                m[leg][2]->setTorque(-tau);
            }
            else m[leg][2]->setPosition(- msg-> q_des_knee[leg] - ang_bias);
        }
    }

    iter ++;
}
void webotsController::run(void) {
    // initiate
    initNode();

    // main loop
    while (robot->step(timeStep) != -1) {
        updateLegData();
        updateCheaterState();
        updateImuData();
        // updateTouchForce();
        ros::spinOnce();
    };
}

int main(int argc, char** argv) {
    ros::init(argc, argv, "Webots_Controller", ros::init_options::AnonymousName);
    ros::NodeHandle n;
    webotsController controller(&n);
    controller.run();
}
